{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "49eaa538",
   "metadata": {},
   "source": [
    "# Tutorial 06, case 3b: advection diffusion reaction control problem with Neumann control\n",
    "\n",
    "In this tutorial we solve the optimal control problem\n",
    "\n",
    "$$\\min J(y, u) = \\frac{1}{2} \\int_{\\Omega_3} (y - y_d)^2 dx + \\frac{\\alpha}{2} \\int_{\\Gamma_2} u^2 ds$$\n",
    "s.t.\n",
    "$$\\begin{cases}\n",
    "- \\epsilon \\Delta y + \\beta \\cdot \\nabla y + \\sigma y = f      & \\text{in } \\Omega\\\\\n",
    "                                                    y = g      & \\text{on } \\Gamma_1\\\\\n",
    "                                \\epsilon \\partial_n y = u      & \\text{on } \\Gamma_2\\\\\n",
    "                                \\epsilon \\partial_n y = 0      & \\text{on } \\Gamma_3\\\\\n",
    "\\end{cases}$$\n",
    "\n",
    "where\n",
    "$$\\begin{align*}\n",
    "& \\Omega               & \\text{domain}\\\\\n",
    "& u \\in L^2(\\Gamma_2)  & \\text{control variable}\\\\\n",
    "& y \\in H^1(\\Omega)    & \\text{state variable}\\\\\n",
    "& \\alpha > 0           & \\text{penalization parameter}\\\\\n",
    "& y_d                  & \\text{desired state}\\\\\n",
    "& f                    & \\text{forcing term}\\\\\n",
    "& g                    & \\text{nonhomogeneous Dirichlet BC}\\\\\n",
    "\\end{align*}$$\n",
    "using an adjoint formulation solved by a one shot approach.\n",
    "\n",
    "The test case is from section 5.3 of\n",
    "```\n",
    "F. Negri, G. Rozza, A. Manzoni and A. Quarteroni. Reduced Basis Method for Parametrized Elliptic Optimal Control Problems. SIAM Journal on Scientific Computing, 35(5): A2316-A2340, 2013.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9357a530",
   "metadata": {},
   "outputs": [],
   "source": [
    "import typing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9b45d44",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfinx.fem\n",
    "import dolfinx.fem.petsc\n",
    "import dolfinx.io\n",
    "import dolfinx.mesh\n",
    "import gmsh\n",
    "import mpi4py.MPI\n",
    "import numpy as np\n",
    "import petsc4py.PETSc\n",
    "import ufl\n",
    "import viskex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cd8a3d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiphenicsx.fem\n",
    "import multiphenicsx.fem.petsc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c2b0812",
   "metadata": {},
   "source": [
    "### Geometrical parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17e56000",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh_size = 0.05"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94bc00e8",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f67db316",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GenerateRectangleLines:\n",
    "    \"\"\"Generate a rectangle.\"\"\"\n",
    "\n",
    "    def __init__(self) -> None:\n",
    "        self.points: dict[tuple[float, float], int] = dict()\n",
    "        self.lines: dict[tuple[int, int], int] = dict()\n",
    "\n",
    "    def add_point(self, x: float, y: float) -> int:\n",
    "        \"\"\"Add a point to gmsh, if not present already.\"\"\"\n",
    "        key = (x, y)\n",
    "        try:\n",
    "            return self.points[key]\n",
    "        except KeyError:\n",
    "            p = gmsh.model.geo.addPoint(x, y, 0.0, mesh_size)\n",
    "            self.points[key] = p\n",
    "            return p  # type: ignore[no-any-return]\n",
    "\n",
    "    def add_line(self, p0: int, p1: int) -> int:\n",
    "        \"\"\"Add a line to gmsh, if not present already.\"\"\"\n",
    "        try:\n",
    "            return self.lines[(p0, p1)]\n",
    "        except KeyError:\n",
    "            l01 = gmsh.model.geo.addLine(p0, p1)\n",
    "            self.lines[(p0, p1)] = l01\n",
    "            self.lines[(p1, p0)] = -l01\n",
    "            return l01  # type: ignore[no-any-return]\n",
    "\n",
    "    def __call__(\n",
    "        self, x_min: float, x_max: float, y_min: float, y_max: typing.Union[float, list[float]]\n",
    "    ) -> tuple[int, list[int], int, int]:\n",
    "        \"\"\"Add points and lines that define a rectangle with the provided coordinates.\"\"\"\n",
    "        p0 = self.add_point(x_min, y_min)\n",
    "        p1 = self.add_point(x_max, y_min)\n",
    "        if isinstance(y_max, list):\n",
    "            p2 = [self.add_point(x_max, y) for y in y_max]\n",
    "            p3 = self.add_point(x_min, y_max[-1])\n",
    "        else:\n",
    "            p2 = [self.add_point(x_max, y_max)]\n",
    "            p3 = self.add_point(x_min, y_max)\n",
    "        l0 = self.add_line(p0, p1)\n",
    "        p1_p2 = [p1, *p2]\n",
    "        l1 = [self.add_line(p1_p2[i], p1_p2[i + 1]) for i in range(len(p2))]\n",
    "        l2 = self.add_line(p2[-1], p3)\n",
    "        l3 = self.add_line(p3, p0)\n",
    "        return (l0, l1, l2, l3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca6d3b9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmsh.initialize()\n",
    "gmsh.model.add(\"mesh\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71e87b76",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_rectangle_lines = GenerateRectangleLines()\n",
    "[l0, l1, l2, l3] = generate_rectangle_lines(0.0, 1.0, 0.0, [0.3, 0.7, 1.0])\n",
    "[l4, l5, l6, _] = generate_rectangle_lines(1.0, 3.0, 0.0, 0.3)\n",
    "[_, l7, l8, _] = generate_rectangle_lines(1.0, 3.0, 0.3, 0.7)\n",
    "[_, l9, l10, _] = generate_rectangle_lines(1.0, 3.0, 0.7, 1.0)\n",
    "line_loop_rectangle_left = gmsh.model.geo.addCurveLoop([l0, l1[0], l1[1], l1[2], l2, l3])\n",
    "line_loop_rectangle_right_bottom = gmsh.model.geo.addCurveLoop([l4, l5[0], l6, -l1[0]])\n",
    "line_loop_rectangle_right_middle = gmsh.model.geo.addCurveLoop([-l6, l7[0], l8, -l1[1]])\n",
    "line_loop_rectangle_right_top = gmsh.model.geo.addCurveLoop([-l8, l9[0], l10, -l1[2]])\n",
    "rectangle_left = gmsh.model.geo.addPlaneSurface([line_loop_rectangle_left])\n",
    "rectangle_right_bottom = gmsh.model.geo.addPlaneSurface([line_loop_rectangle_right_bottom])\n",
    "rectangle_right_middle = gmsh.model.geo.addPlaneSurface([line_loop_rectangle_right_middle])\n",
    "rectangle_right_top = gmsh.model.geo.addPlaneSurface([line_loop_rectangle_right_top])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c8899e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmsh.model.geo.synchronize()\n",
    "gmsh.model.addPhysicalGroup(1, [l0, l2, l3], 1)\n",
    "gmsh.model.addPhysicalGroup(1, [l4, l10], 2)\n",
    "gmsh.model.addPhysicalGroup(1, [l5[0], l7[0], l9[0]], 3)\n",
    "gmsh.model.addPhysicalGroup(2, [rectangle_left], 1)\n",
    "gmsh.model.addPhysicalGroup(2, [rectangle_right_bottom, rectangle_right_top], 3)\n",
    "gmsh.model.addPhysicalGroup(2, [rectangle_right_middle], 2)\n",
    "gmsh.model.mesh.generate(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85f21c7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh, subdomains, boundaries = dolfinx.io.gmshio.model_to_mesh(\n",
    "    gmsh.model, comm=mpi4py.MPI.COMM_WORLD, rank=0, gdim=2)\n",
    "gmsh.finalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cd851ef-d028-4545-b2cf-7457d3fbc722",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create connectivities required by the rest of the code\n",
    "mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aebf031",
   "metadata": {},
   "outputs": [],
   "source": [
    "boundaries_1 = boundaries.indices[boundaries.values == 1]\n",
    "boundaries_2 = boundaries.indices[boundaries.values == 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8bdaa52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define associated measures\n",
    "dx = ufl.Measure(\"dx\", subdomain_data=subdomains)\n",
    "ds = ufl.Measure(\"ds\", subdomain_data=boundaries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "952ab416",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93045d0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh_tags(mesh, subdomains, \"subdomains\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8759a1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh_tags(mesh, boundaries, \"boundaries\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2e81fee",
   "metadata": {},
   "source": [
    "### Function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc46de18",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 2))\n",
    "U = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 2))\n",
    "Q = Y.clone()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d73b0f4",
   "metadata": {},
   "source": [
    "### Restrictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9257414d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dofs_Y = np.arange(0, Y.dofmap.index_map.size_local + Y.dofmap.index_map.num_ghosts)\n",
    "dofs_U = dolfinx.fem.locate_dofs_topological(U, boundaries.dim, boundaries_2)\n",
    "dofs_Q = dofs_Y\n",
    "restriction_Y = multiphenicsx.fem.DofMapRestriction(Y.dofmap, dofs_Y)\n",
    "restriction_U = multiphenicsx.fem.DofMapRestriction(U.dofmap, dofs_U)\n",
    "restriction_Q = multiphenicsx.fem.DofMapRestriction(Q.dofmap, dofs_Q)\n",
    "restriction = [restriction_Y, restriction_U, restriction_Q]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09197c6e",
   "metadata": {},
   "source": [
    "### Trial and test functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f8d9e4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "(y, u, p) = (ufl.TrialFunction(Y), ufl.TrialFunction(U), ufl.TrialFunction(Q))\n",
    "(z, v, q) = (ufl.TestFunction(Y), ufl.TestFunction(U), ufl.TestFunction(Q))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "193f5866",
   "metadata": {},
   "source": [
    " ### Problem data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0e07e14",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 0.07\n",
    "y_d = 2.5\n",
    "epsilon = 1. / 12.\n",
    "x = ufl.SpatialCoordinate(mesh)\n",
    "beta = ufl.as_vector((x[1] * (1 - x[1]), 0))\n",
    "sigma = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0))\n",
    "ff = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0))\n",
    "bc0 = petsc4py.PETSc.ScalarType(0)\n",
    "bc1 = petsc4py.PETSc.ScalarType(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4775a36a",
   "metadata": {},
   "source": [
    "### Optimality conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "901987fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_operator = (epsilon * ufl.inner(ufl.grad(y), ufl.grad(q)) * dx\n",
    "                  + ufl.inner(ufl.dot(beta, ufl.grad(y)), q) * dx + sigma * ufl.inner(y, q) * dx)\n",
    "adjoint_operator = (epsilon * ufl.inner(ufl.grad(p), ufl.grad(z)) * dx\n",
    "                    - ufl.inner(ufl.dot(beta, ufl.grad(p)), z) * dx + sigma * ufl.inner(p, z) * dx)\n",
    "a = [[ufl.inner(y, z) * dx(3), None, adjoint_operator],\n",
    "     [None, alpha * ufl.inner(u, v) * ds(2), - ufl.inner(p, v) * ds(2)],\n",
    "     [state_operator, - ufl.inner(u, q) * ds(2), None]]\n",
    "f = [ufl.inner(y_d, z) * dx(3),\n",
    "     None,\n",
    "     ufl.inner(ff, q) * dx]\n",
    "a[0][0] += dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)) * ufl.inner(y, z) * dx\n",
    "a[2][2] = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)) * ufl.inner(p, q) * dx\n",
    "f[1] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), v) * dx\n",
    "a_cpp = dolfinx.fem.form(a)\n",
    "f_cpp = dolfinx.fem.form(f)\n",
    "bdofs_Y_1 = dolfinx.fem.locate_dofs_topological(Y, mesh.topology.dim - 1, boundaries_1)\n",
    "bdofs_Q_1 = dolfinx.fem.locate_dofs_topological(Q, mesh.topology.dim - 1, boundaries_1)\n",
    "bc = [dolfinx.fem.dirichletbc(bc1, bdofs_Y_1, Y),\n",
    "      dolfinx.fem.dirichletbc(bc0, bdofs_Q_1, Q)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32c54437",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b442e85",
   "metadata": {},
   "outputs": [],
   "source": [
    "(y, u, p) = (dolfinx.fem.Function(Y), dolfinx.fem.Function(U), dolfinx.fem.Function(Q))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16ad20e9",
   "metadata": {},
   "source": [
    "### Cost functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3842b8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "J = 0.5 * ufl.inner(y - y_d, y - y_d) * dx(3) + 0.5 * alpha * ufl.inner(u, u) * ds(2)\n",
    "J_cpp = dolfinx.fem.form(J)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fcededa",
   "metadata": {},
   "source": [
    "### Uncontrolled functional value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76070dd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract state forms from the optimality conditions\n",
    "a_state = ufl.replace(a[2][0], {q: z})\n",
    "f_state = ufl.replace(f[2], {q: z})\n",
    "a_state_cpp = dolfinx.fem.form(a_state)\n",
    "f_state_cpp = dolfinx.fem.form(f_state)\n",
    "bc_state = [bc[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d87930ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the linear system for the state\n",
    "A_state = dolfinx.fem.petsc.assemble_matrix(a_state_cpp, bcs=bc_state)\n",
    "A_state.assemble()\n",
    "F_state = dolfinx.fem.petsc.assemble_vector(f_state_cpp)\n",
    "dolfinx.fem.petsc.apply_lifting(F_state, [a_state_cpp], [bc_state])\n",
    "F_state.ghostUpdate(addv=petsc4py.PETSc.InsertMode.ADD, mode=petsc4py.PETSc.ScatterMode.REVERSE)\n",
    "dolfinx.fem.petsc.set_bc(F_state, bc_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad6b44a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A_state)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F_state, y.x.petsc_vec)\n",
    "y.x.petsc_vec.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a93740e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "J_uncontrolled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)\n",
    "print(\"Uncontrolled J =\", J_uncontrolled)\n",
    "assert np.isclose(J_uncontrolled, 1.35)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56d89294",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(y, \"uncontrolled state\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b438f45",
   "metadata": {},
   "source": [
    "### Optimal control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b3047e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system for the optimality conditions\n",
    "A = multiphenicsx.fem.petsc.assemble_matrix_block(a_cpp, bcs=bc, restriction=(restriction, restriction))\n",
    "A.assemble()\n",
    "F = multiphenicsx.fem.petsc.assemble_vector_block(f_cpp, a_cpp, bcs=bc, restriction=restriction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0c04801",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "yup = multiphenicsx.fem.petsc.create_vector_block(f_cpp, restriction=restriction)\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, yup)\n",
    "yup.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ee8511e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(yup, [Y.dofmap, U.dofmap, Q.dofmap], restriction) as yup_wrapper:\n",
    "    for yup_wrapper_local, component in zip(yup_wrapper, (y, u, p)):\n",
    "        with component.x.petsc_vec.localForm() as component_local:\n",
    "            component_local[:] = yup_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aff05068",
   "metadata": {},
   "outputs": [],
   "source": [
    "J_controlled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)\n",
    "print(\"Optimal J =\", J_controlled)\n",
    "assert np.isclose(J_controlled, 0.035933676)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a968ea8",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(y, \"state\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbef7fa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u, \"control\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f4a20e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(p, \"adjoint\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
